package com.gitee.wsl.coroutines.channels

import kotlinx.atomicfu.atomic
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.Deferred
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock

class MultiChannel<T> {

    private var isClosed: Boolean = false
    private var closeException: Throwable? = null
    private val lock = Mutex()
    private val pending = mutableListOf<Pair<String, CompletableDeferred<T>>>()
    private val id = atomic(1)

    private fun checkClosed() {
        if (isClosed) {
            throw IllegalStateException("MultiChannel ($this) is closed", closeException)
        }
    }

    suspend fun send(id: String, response: T) {
        checkClosed()
        lock.withLock {
            val hasPending = pending.consume(matcher = { it.first == id }) { (_, pendingItem) ->
                pendingItem.complete(response)
            }
            if (!hasPending) {
                error("No pending receiver for $id and $response")
            }
        }
    }

    suspend fun allocateReceive(): Pair<Int, Deferred<T>> {
        checkClosed()
        val id = this.id.getAndIncrement()
        val completable = CompletableDeferred<T>()
        lock.withLock {
            pending.add(id.toString() to completable)
        }
        return id to completable
    }

    suspend fun close(t: CancellationException? = null) {
        lock.withLock {
            isClosed = true
            closeException = t
            pending.forEach {
                it.second.completeExceptionally(t ?: CancellationException("Closing MultiChannel"))
            }
        }
    }
}

internal inline fun <T> MutableList<T>.consume(
    crossinline matcher: (T) -> Boolean,
    crossinline consumer: (T) -> Unit
): Boolean {
    return removeAll {
        if (matcher(it)) {
            consumer(it)
            true
        } else {
            false
        }
    }
}
